import numpy as np
from sklearn.datasets import load_wine
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.preprocessing import StandardScaler


def classification(train_feature, train_label, test_feature):
    '''
    对test_feature进行红酒分类
    :param train_feature: 训练集数据，类型为ndarray
    :param train_label: 训练集标签，类型为ndarray
    :param test_feature: 测试集数据，类型为ndarray
    :return: 测试集数据的分类结果
    '''
    md = KNeighborsClassifier(metric="mahalanobis", metric_params={
                              "V": np.cov(train_feature.T)})
    md.fit(train_feature, train_label)
    return md.predict(test_feature)


if __name__ == '__main__':
    wine_dataset = load_wine()
    X, Y = wine_dataset.data, wine_dataset.target
    X_train, X_test, Y_train, Y_test = train_test_split(X, Y, random_state=233)
    pred = classification(X_train, Y_train, X_test)
    acc = accuracy_score(Y_test, pred)
    if acc > 0.92:
        print('你的分类准确率高于0.92')
    else:
        print('你的分类准确率为%.6f' % acc)
